import math
import time
import sys
import logging
import os
import csv
from datetime import datetime
from typing import Dict, Tuple
from getDCSdata import get_dcs_data
from vjoy_control_new import (
    activate_dcs_window,
    throttle_down,  set_throttle,
    foot_left, foot_right,
    airbrake_on_control, airbrake_off_control,
    wheel_control, flapsland_control, gearson_control,
    pitch_down
)

# 全局配置常量
RUNWAY_CONFIG = {
    'LATITUDE': 41.929631,  # 跑道纬度坐标
    'LONGITUDE': 41.86325,  # 跑道经度坐标
    'WIDTH': 50,  # 跑道宽度（米）
    'MAX_REWARD': 100,  # 最大奖励值
    'REWARD_RATE': 0.5,  # 奖励计算系数
    'CONTROL_INTERVAL': 0.3,  # 控制循环间隔（秒）
    'WINDOW_ACTIVATION_INTERVAL': 5,  # DCS窗口激活间隔（秒）
    'MAX_RETRY_COUNT': 3,  # 最大重试次数
    'RETRY_DELAY': 2  # 重试延迟（秒）
}

# 减速阶段定义 (单位: m/s)
DECEL_STAGES = {
    'high_speed': (70, 200),  # 高速阶段
    'mid_speed': (30, 70),  # 中速阶段
    'low_speed': (5, 30),  # 低速阶段
    'stopped': (0, 5)  # 停止状态
}

# 确保结果目录存在
RESULTS_DIR = "deceleration_results"
os.makedirs(RESULTS_DIR, exist_ok=True)

# 日志系统配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(os.path.join(RESULTS_DIR, 'runway_deceleration.log')),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class RunwayDecelerationController:
    """优化后的跑道减速控制器（vJoy版）"""

    def __init__(self):
        """初始化控制器状态"""
        self.target_heading = None
        self.last_operation_time = time.time()
        self.last_window_activation = 0
        self.last_ias = 0
        self.last_time = time.time()
        self.acceleration_threshold = 1.5
        self.last_status = None
        self.total_reward = 0
        self.on_runway = False
        self.reward_active = False
        self.smoothed_accel = 0
        self.decel_stage = None
        self.brake_pulse_count = 0
        self.last_brake_time = 0
        self.control_history = []

        # 初始化数据记录文件
        self.data_file = os.path.join(RESULTS_DIR, f"deceleration_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv")
        self._init_data_file()

    def _init_data_file(self):
        """初始化CSV数据文件"""
        with open(self.data_file, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow([
                'timestamp', 'speed', 'acceleration', 'decel_stage',
                'throttle_action', 'brake_action', 'airbrake_action',
                'rudder_action', 'pitch_action', 'reward', 'total_reward'
            ])

    def _record_data(self, data: Dict):
        """记录控制数据"""
        with open(self.data_file, 'a', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow([
                datetime.now().strftime('%H:%M:%S.%f'),
                data.get('IAS', 0),
                self.smoothed_accel,
                self.decel_stage,
                data.get('Throttle', 0),
                data.get('WheelBrake', 0),
                data.get('AirBrake', 0),
                data.get('Rudder', 0),
                data.get('Pitch', 0),
                self.last_reward if hasattr(self, 'last_reward') else 0,
                self.total_reward
            ])

    def activate_simulator_window(self) -> bool:
        """激活模拟游戏窗口"""
        try:
            if time.time() - self.last_window_activation > RUNWAY_CONFIG['WINDOW_ACTIVATION_INTERVAL']:
                activate_dcs_window()
                self.last_window_activation = time.time()
                logger.info("模拟窗口激活成功")
            return True
        except Exception as error:
            logger.error(f"模拟窗口激活失败: {error}")
            return False

    def get_current_status(self):
        """获取并验证飞行状态数据"""
        try:
            status = get_dcs_data()
            if status is None:
                print("收到空的状态数据")
                return None

            # 设置跑道方向（以当前航向作为初始跑道方向）
            if self.target_heading is None:
                self.target_heading = status['heading_deg']
                print(f"初始目标航向设置为: {self.target_heading:.1f}°")
            status['RunwayHeading'] = self.target_heading

            # 验证必要字段
            necessary_fields = ['IAS', 'LLA_balt', 'heading_deg', 'Gears', 'Flaps', 'AirBrake', 'VV_0']
            for field in necessary_fields:
                if field not in status:
                    print(f"状态数据缺少必要字段: {field}")
                    return None

            # 调试输出
            print(
                f"状态更新 | 速度: {status['IAS']:.1f} m/s | 高度: {status['LLA_balt']:.1f} m | 航向: {status['heading_deg']:.1f}°| 起落架: {status['Gears']:.1f}  | 襟翼: {status['Flaps']:.1f} | 垂直速度: {status['VV_0']:.1f} ")
            return status

        except Exception as error:
            print(f"获取状态数据时发生异常: {str(error)}")
            return None

    def check_runway_status(self, current_status: Dict) -> bool:
        """检查是否进入跑道状态"""
        vertical_speed = current_status.get('VV_0', 0)
        runway_conditions = (
                current_status['LLA_balt'] < 300 and
                abs(vertical_speed) < 10 and  # 垂直速度小于10m/s
                current_status['IAS'] < 300 and
                current_status.get('Gears', 0) > 0.8
        )
        # 添加调试日志
        logger.debug(
            f"跑道条件检查 | 高度:{current_status['LLA_balt']:.1f}m(<300) "
            f"垂直速度:{vertical_speed:.1f}m/s(<10) "
            f"速度:{current_status['IAS']:.1f}m/s(<300) "
            f"起落架:{current_status.get('Gears', 0):.1f}(>0.8) "
            f"结果:{runway_conditions}"
        )

        if runway_conditions:
            if not self.on_runway:
                self._enter_runway(current_status)
            return True
        else:
            if self.on_runway:
                self._exit_runway()
            return False

    def _enter_runway(self, status: Dict):
        """进入跑道区域处理"""
        self.on_runway = True
        self.last_ias = status['IAS']
        self.reward_active = True

        # 进入跑道时的标准配置
        #  flapsland_control() # 襟翼全放
        gearson_control()  # 确保起落架放下

        logger.info(f"进入跑道区域 | 速度: {status['IAS']:.1f}m/s | 高度: {status['LLA_balt']:.1f}m")

    def _exit_runway(self):
        """离开跑道区域处理"""
        self.on_runway = False
        self.reward_active = False
        logger.info(f"离开跑道区域 | 累计奖励: {self.total_reward:.1f}/{RUNWAY_CONFIG['MAX_REWARD']}")

    def _determine_decel_stage(self, speed: float) -> str:
        """确定当前减速阶段"""
        for stage, (min_speed, max_speed) in DECEL_STAGES.items():
            if min_speed <= speed < max_speed:
                return stage
        return 'unknown'

    def _control_brakes(self, current_status: Dict):
        """智能刹车控制"""
        current_time = time.time()
        airspeed = current_status["IAS"]
        acceleration = self.calculate_acceleration(airspeed)

        # 根据速度阶段采用不同的刹车策略
        if self.decel_stage == 'high_speed':
            if current_status.get("AirBrake", 0) < 0.5:
                airbrake_on_control()  # 只在减速板未完全展开时激活

            # 脉冲式使用刹车，间隔时间根据速度动态调整
            brake_interval = max(0.5,1.5 - airspeed / 200)  # 速度越高，间隔越短
            if current_time - self.last_brake_time > brake_interval:
                wheel_control()
                self.last_brake_time = current_time

        elif self.decel_stage == 'mid_speed':
            # 中速阶段组合使用减速板和刹车
            if acceleration > 0.5:  # 减速不足
                airbrake_on_control()
                wheel_control()
            else:
                if current_time - self.last_brake_time > 0.5:
                    wheel_control()
                    self.last_brake_time = current_time

        elif self.decel_stage == 'low_speed':
            # 低速阶段主要使用刹车
            if current_time - self.last_brake_time > 0.3:
                wheel_control()
                self.last_brake_time = current_time
            # 关闭减速板防止过度减速
            if current_status["AirBrake"] > 0.5:
                airbrake_off_control()

    def _control_rudder(self, current_status: Dict):
        """方向舵控制保持航向"""
        heading_error = abs((current_status["heading_deg"] - self.target_heading + 180) % 360 - 180)

        # 只有当误差足够大时才修正，且修正量更平缓
        if heading_error > 3:  # 降低阈值
            correction_amount = min(0.7, heading_error / 40)  # 减小最大修正量
            duration = min(0.3, heading_error / 80)  # 缩短修正时间

            # 根据误差方向决定修正方向
            if (current_status["heading_deg"] - self.target_heading + 180) % 360 - 180 > 0:
                foot_right(duration=correction_amount)  # 使用duration参数控制幅度
            else:
                foot_left(duration=correction_amount)

            logger.debug(f"航向修正 | 误差: {heading_error:.1f}° | 修正量: {correction_amount:.2f}")

    def _control_pitch(self, current_status: Dict):
        """俯仰控制保持前轮接地"""
        if self.decel_stage in ['high_speed', 'mid_speed']:
            # 中高速阶段保持轻微拉杆
            pitch_down(duration=0.1)  # 使用duration参数控制幅度
        elif self.decel_stage == 'low_speed':
            # 低速阶段减小俯仰角
            pitch_down(duration=0.2)

    def _control_throttle(self, current_status: Dict):
        """油门控制"""
        if current_status["IAS"] > 30:
            set_throttle(0.0)  # 高速阶段油门收到底
        elif current_status["IAS"] > 10:
            current_throttle = max(0.0, current_status.get("Throttle", 0))
            if current_throttle > 0.1:
                throttle_down(0.1)  # 逐步减少油门
        else:
            set_throttle(0.0)  # 低速阶段确保油门关闭

    def perform_deceleration_control(self, current_status: Dict) -> bool:
        """执行减速控制逻辑"""
        current_time = time.time()
        if current_time - self.last_operation_time < RUNWAY_CONFIG['CONTROL_INTERVAL']:
            return False

        self.last_operation_time = current_time
        airspeed = current_status["IAS"]
        self.decel_stage = self._determine_decel_stage(airspeed)

        # 1. 执行各控制子系统
        self._control_brakes(current_status)
        self._control_rudder(current_status)
        self._control_pitch(current_status)
        self._control_throttle(current_status)

        # 2. 奖励计算
        if self.last_status and self.reward_active:
            reward, details = self.calculate_runway_deceleration_reward(self.last_status, current_status)
            self.last_reward = reward
            self.total_reward += reward
            logger.info(
                f"减速阶段: {self.decel_stage:10s} | "
                f"速度: {airspeed:5.1f}m/s | "
                f"本次奖励: +{reward:4.2f} | "
                f"累计: {self.total_reward:5.1f}/{RUNWAY_CONFIG['MAX_REWARD']}"
            )

        # 3. 记录控制历史（用于调试）
        self.control_history.append({
            'time': current_time,
            'speed': airspeed,
            'stage': self.decel_stage,
            'actions': {
                'brake': current_status.get("WheelBrake", 0),
                'airbrake': current_status.get("AirBrake", 0),
                'rudder': current_status.get("Rudder", 0)
            }
        })

        # 4. 记录数据
        self._record_data(current_status)

        # 5. 停止条件检查
        if airspeed < 5 and current_status['LLA_balt'] < 5:
            logger.info(f"飞机完全停止！最终奖励: {self.total_reward:.1f}/{RUNWAY_CONFIG['MAX_REWARD']}")
            self._generate_deceleration_report()
            return True

        self.last_status = current_status.copy()
        return False

    def calculate_acceleration(self, current_ias: float) -> float:
        """计算平滑后的加速度"""
        current_time = time.time()
        time_diff = current_time - self.last_time

        if time_diff > 0:
            raw_accel = (current_ias - self.last_ias) / time_diff
            self.smoothed_accel = 0.7 * self.smoothed_accel + 0.3 * raw_accel
            self.last_ias = current_ias
            self.last_time = current_time

        return self.smoothed_accel

    def calculate_runway_deceleration_reward(self, prev_state: Dict, current_state: Dict) -> Tuple[float, Dict]:
        """计算跑道减速奖励"""
        ias = current_state["IAS"]
        prev_ias = prev_state["IAS"]
        heading_error = abs((current_state["heading_deg"] - self.target_heading + 180) % 360 - 180)
        distance = current_state.get("RunwayDistance", RUNWAY_CONFIG['WIDTH'] * 2)

        # 减速奖励 (基于速度变化)
        speed_diff = max(0, prev_ias - ias)
        decel_reward = min(speed_diff * RUNWAY_CONFIG['MAX_REWARD'], 0.4 * RUNWAY_CONFIG['MAX_REWARD'])

        # 航向保持奖励
        alignment_reward = 0.3 * RUNWAY_CONFIG['MAX_REWARD'] * math.exp(-heading_error / 5.0)

        # 跑道中线保持奖励
        distance_penalty = min(distance / RUNWAY_CONFIG['WIDTH'], 1)
        distance_reward = 0.3 * RUNWAY_CONFIG['MAX_REWARD'] * (1 - distance_penalty)

        # 阶段加成
        stage_bonus = {
            'high_speed': 1.2,
            'mid_speed': 1.0,
            'low_speed': 0.8
        }.get(self.decel_stage, 1.0)

        total_reward = max(0, (decel_reward + alignment_reward + distance_reward) * stage_bonus)
        return min(total_reward, RUNWAY_CONFIG['MAX_REWARD'] - self.total_reward), {
            "decel_reward": decel_reward,
            "alignment_reward": alignment_reward,
            "distance_reward": distance_reward,
            "stage_bonus": stage_bonus
        }

    def _generate_deceleration_report(self):
        """生成减速过程报告"""
        if not self.control_history:
            logger.info("没有控制历史数据，无法生成报告")
            return

        try:
            # 创建报告文件
            report_file = os.path.join(RESULTS_DIR,
                f"deceleration_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")

            with open(report_file, 'w') as f:
                f.write("=== 跑道减速过程报告 ===\n\n")
                f.write(f"减速完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"最终奖励: {self.total_reward:.1f}/{RUNWAY_CONFIG['MAX_REWARD']}\n\n")

                f.write("减速阶段分析:\n")
                # 计算各阶段耗时
                stage_times = {}
                prev_time = self.control_history[0]['time']
                prev_stage = self.control_history[0]['stage']

                for record in self.control_history[1:]:
                    current_stage = record['stage']
                    if current_stage != prev_stage:
                        if prev_stage not in stage_times:
                            stage_times[prev_stage] = 0
                        stage_times[prev_stage] += record['time'] - prev_time
                        prev_stage = current_stage
                        prev_time = record['time']

                # 添加最后一个阶段
                if prev_stage not in stage_times:
                    stage_times[prev_stage] = 0
                stage_times[prev_stage] += self.control_history[-1]['time'] - prev_time

                # 写入阶段耗时
                for stage, duration in stage_times.items():
                    f.write(f"  {stage}: {duration:.1f}秒\n")

                f.write("\n控制统计:\n")
                # 计算各控制动作使用频率
                brake_count = sum(1 for r in self.control_history if r['actions']['brake'] > 0.5)
                airbrake_count = sum(1 for r in self.control_history if r['actions']['airbrake'] > 0.5)
                rudder_count = sum(1 for r in self.control_history if r['actions']['rudder'] > 0.5)

                f.write(f"  刹车使用次数: {brake_count}\n")
                f.write(f"  减速板使用次数: {airbrake_count}\n")
                f.write(f"  方向舵修正次数: {rudder_count}\n")

                f.write("\n改进建议:\n")
                if self.total_reward < RUNWAY_CONFIG['MAX_REWARD'] * 0.7:
                    f.write("  - 优化刹车时机，避免过早或过晚使用\n")
                    f.write("  - 保持更好的跑道中心线对齐\n")
                    f.write("  - 在高速阶段更有效地使用减速板\n")
                else:
                    f.write("  - 减速过程表现良好，继续保持\n")

            logger.info(f"减速报告已保存至 {report_file}")

        except Exception as e:
            logger.error(f"生成减速报告时出错: {str(e)}")

    def run(self) -> bool:
        """主控制循环"""
        logger.info("==== 跑道减速控制系统启动 ====")
        if not self.activate_simulator_window():
            return False

        try:
            while True:
                current_status = self.get_current_status()
                if current_status is None:
                    time.sleep(1)
                    continue

                if self.check_runway_status(current_status):
                    if current_status['IAS'] > 5:
                        if self.perform_deceleration_control(current_status):
                            break
                    else:
                        logger.info("飞机已完全停止")
                        break

                # 定期激活窗口防止失去焦点
                if time.time() - self.last_window_activation > RUNWAY_CONFIG['WINDOW_ACTIVATION_INTERVAL']:
                    self.activate_simulator_window()

                time.sleep(RUNWAY_CONFIG['CONTROL_INTERVAL'])

        except KeyboardInterrupt:
            logger.info("用户中断控制")
        except Exception as error:
            logger.error(f"控制循环出错: {error}")
            return False
        finally:
            logger.info("==== 系统关闭 ====")
        return True


if __name__ == "__main__":
    controller = RunwayDecelerationController()
    try:
        success = controller.run()
        sys.exit(0 if success else 1)
    except Exception as error:
        logger.error(f"未处理的异常: {error}")
        sys.exit(1)